"""
esp01s驱动模块
最后更新日期: 2024/12/21
说明: 本模块仅提供针对esp01s的操作
"""

import utime


TIMEOUT_DEFAULT = 2000  #默认的超时时间

# 通用指令
CMD_ONLINE_TEST = 'AT'  # 测试命令
CMD_RESET = 'AT+RST'  # 复位
CMD_RESTORE = 'AT+RESTORE'  # 复位
CMD_SYSINFO  = 'AT+GMR'  # 查看系统版本
CMD_BAUDRATE  = 'AT+UART={baudrate},8,1,0,0'  # 设置波特率
CMD_SET_CWMODE = 'AT+CWMODE={mode}'  # 设置模式
CMD_GET_CWMODE = 'AT+CWMODE?'  # 查询模式
CMD_GET_MYSELF_IP_MAC  = 'AT+CIFSR'  # 本机ip和mac

# 不同工作模式的指令
CMD_EXIT_SEND_MODE = "+++"  # 退出发送模式
CMD_CONNECT_TO_AP  = 'AT+CWJAP="{ssid}","{pwd}"'  # 连接ap
CMD_GET_CWJAP  = "AT+CWJAP?"  # 获取连接的ap信息
CMD_  = 'AT+'  #
CMD_  = 'AT+'  #
CMD_  = 'AT+'  #
CMD_  = 'AT+'  # 

# MQTT固件专用指令
CMD_SET_NTP = 'AT+CIPSNTPCFG=1,8,{ntp_server}'  # 设置ntp设置
CMD_GET_NTP = 'AT+CIPSNTPCFG?'  # 获取ntp设置
CMD_GET_TIME  = 'AT+CIPSNTPTIME?'  # 获取时间
CMD_SET_MQTTUSERCFG = 'AT+MQTTUSERCFG=0,1,"{client_id}","","",0,0,""'  #mqtt用户设置,参数为私钥
CMD_SET_MQTTCONNCFG = 'AT+MQTTCONNCFG=0,5,0,"{topic}","下线了",0,0'  # mqtt连接设置, 遗嘱设置
CMD_SET_MQTTCONN = 'AT+MQTTCONN=0,"{host}",{port},1'  # mqtt连接设置, 参数为主机地址和端口号
CMD_GET_MQTTCONN = 'AT+MQTTCONN?'
CMD_SET_MQTTSUB = 'AT+MQTTSUB=0,"{topic}",0'  # mqtt订阅主题,参数为主题
CMD_GET_MQTTSUB = 'AT+MQTTSUB?'
CMD_SET_MQTTUNSUB = 'AT+MQTTUNSUB=0,"{topic}"'  # mqtt取消订阅,参数为主题
CMD_SET_MQTTPUB = 'AT+MQTTPUB=0,"{topic}","{data}",0,0'  # mqtt发布,参数为主题和消息内容
CMD_SET_MQTTCLEAN = 'AT+MQTTCLEAN=0'  # mqtt断开连接
CMD__MQTTSUBRECV = 'MQTTSUBRECV'  # 接收订阅消息的标识, 不是用来发送的指令


# 返回的状态
STATUS_OK = "OK"  # 系统返回,正常
STATUS_ERROR = "ERROR"  # 系统返回,错误
STATUS_DEFAULT = "DEFAULT"  # 自定义返回, 默认
STATUS_TIMEOUT = "TIMEOUT"  # 自定义返回,超时接收
STATUS_DECODE_ERROR = "DECODE_ERROR"  # 自定义返回,解码错误
STATUS_CONNECTED = "CONNECTED"  # 自定义返回,连接成功
STATUS_GOT_IP = "WIFI GOT IP"  # 自定义返回,连接成功
STATUS_DISCONNECTED = "DISCONNECTED"  # 自定义返回,连接失败
STATUS_NOT_PASS_MUSTER = 'NOT_PASS_MUSTER'  # 自定义返回,消息不符合要求


# 返回状态,在响应的尾部,用来判断是否接收完成. 键为结束语,值为对应的返回状态
REPONSE_END = {
    b"OK\r\n": STATUS_OK,
    b"ERROR\r\n": STATUS_ERROR,
    b'WIFI CONNECTED\r\n': STATUS_CONNECTED,
    b'WIFI DISCONNECT\r\n': STATUS_DISCONNECTED,
    }





class AT(object):
    def __init__(self,uart):
        self.uart = uart

    def get_response(self, only_return_status=False, timeout=TIMEOUT_DEFAULT):
        """
        等待响应.如果收到尾部信息就停止接收,并将响应内容中的头部和尾部分离
        only_return_status: 只返回是否成功
        timeout: 超时时间,单位:秒
        返回:
        注意: readline也全是按行读取,有时最后的error是分两行.所以不能边读取边去除\r\n
            而且还有\r\r\n出现

        """
        
        # 初始化
        prvMills = utime.ticks_ms()  # 接收响应的起始时间
        response = b""  # 响应的初始值
        return_status = STATUS_DEFAULT  #返回状态
        received_lst=[]  # 存储接收数据的列表,分行
        line_end = True  # 是否行尾,因为readline有时将一行分两次读取,通过该变量组合一行
        
        # 接收数据
        while True:
            # 超时退出
            if utime.ticks_diff(utime.ticks_ms(), prvMills) >= timeout:  
                return_status = STATUS_TIMEOUT
                break
            
            # 循环接收
            if self.uart.any():
                if line_end:
                    received = b""
                
                received += self.uart.readline()
                if received.endswith(b'\n'):  # 确定这是一行的末尾
                    line_end = True
                else:
                    line_end = False
                    continue
                
                # 检查是否是尾部信息,是就主动跳出接收
                if received in REPONSE_END.keys(): 
                    return_status = REPONSE_END[received]
                    break
                else:
                    received_lst.append(received)
        
        # 后处理
        delay = utime.ticks_diff(utime.ticks_ms(), prvMills)  # 计算延迟,单位毫秒
        
        #调试用
#         print("已接收: ", received_lst, return_status)
#         print('--------------------------')
        
        # 解码
        try:
            return_str = b''.join(received_lst[1:]).decode('utf-8').strip()  # 第一行是发送的命令,不需要
        except Exception as e:
            return_status = STATUS_DECODE_ERROR
            return_str = str(e)
        
        # 返回,有的命令只要查看是否ok
        if only_return_status:
            return return_status in [STATUS_OK, STATUS_CONNECTED, STATUS_GOT_IP]
        else:
            return return_status, return_str, delay
    

    # 发送命令
    def send_cmd(self, cmd_str):
        # 补全命令
        if not cmd_str.endswith("\r\n"):
            cmd_str += "\r\n"
        
        cmd_byte = cmd_str.encode('utf-8')
        
        # 调试用
        #print("已发送: ",cmd_byte)
        self.uart.write(cmd_byte)
    
    
    # 是否在线(属性)
    @property
    def online(self):
        self.send_cmd(CMD_ONLINE_TEST)
        return self.get_response(only_return_status=True)

    # cwmode(属性)
    @property
    def cwmode(self):
        self.send_cmd(CMD_GET_CWMODE)
        returned = self.get_response()
        if returned[0] == STATUS_OK and returned[1]:
            cwmode_str = returned[1].split(":")[1]
            return int(cwmode_str)
        else:
            return -1
    
    # 连接的ap信息(属性)
    @property
    def ap_info(self):
        self.send_cmd(CMD_GET_CWJAP)
        returned = self.get_response()
        if returned[0] == STATUS_OK and returned[1] and returned[1] !="No AP":
            cwjap_str = returned[1].split(':"')[1].split(",")
            ssid = cwjap_str[0].replace('"','')
            mac = cwjap_str[1].replace('"','')
            channel = int(cwjap_str[2])
            signal = int(cwjap_str[3])
            return True, '', ssid, mac, channel, signal
        else:
            return False, returned[1], '', '', 0, 0
    
    # 本机ip和mac(属性)
    @property
    def ip_mac(self):
        self.send_cmd(CMD_GET_MYSELF_IP_MAC)
        returned = self.get_response()
        if returned[0] == STATUS_OK and returned[1]:
            ip_mac = returned[1].split("\r\n")
            ip = ip_mac[0].split(",")[1].replace('"','')
            mac = ip_mac[1].split(",")[1].replace('"','')
            return True, '', ip, mac
        else:
            return False, returned[1], '', ''
    
    # ntp设置(属性)
    @property
    def ntp_config(self):
        self.send_cmd(CMD_GET_NTP)
        returned = self.get_response()
        if returned[0] == STATUS_OK and returned[1]:
            ntp_config_str = returned[1].split(':')[1]
            ntp_config_lst = ntp_config_str.split(',')
            is_enable = int(ntp_config_lst[0])
            if is_enable == 1:
                time_zone = int(ntp_config_lst[1])
                ntp_server = ntp_config_lst[2].replace('"','')
            
                return True, '', is_enable, time_zone, ntp_server

        return False, returned[1], 0, 0, ''
        

    # 重置
    def reset(self):
        #self.send_cmd(CMD_EXIT_SEND_MODE)
        self.send_cmd(CMD_RESET)
        return self.get_response(only_return_status=True)
    
    
    # 设置工作模式
    def set_cwmode(self,mode):
        cmd = CMD_SET_CWMODE.format(mode=mode)
        self.send_cmd(cmd)
        return self.get_response(only_return_status=True)
    
    
    # 连接至ap
    def connect_to_ap(self, ssid, password):
        cmd = CMD_CONNECT_TO_AP.format(ssid=ssid, pwd=password)
        self.send_cmd(cmd)
        return self.get_response(only_return_status=True)
        
    
    # 设置ntp服务器(最多支持8个)
    def set_ntp(self, ntp_server_lst):
        ntp_server = ''
        for server in ntp_server_lst:
            ntp_server += (',' if ntp_server else '') + '"' + server + '"'
        cmd = CMD_SET_NTP.format(ntp_server=ntp_server)
        
        self.send_cmd(cmd)
        return self.get_response(only_return_status=True)
    
    
    # 获取网络时间, 为方便校时,直接生成时间元组(注意,最后一个是延迟)
    # 注意: 如果日期是个位数,返回会多个空格('+CIPSNTPTIME:Sun Nov  7 20:17:44 2024'),会导致分隔错误,故要去除列表中的空字符串
    def get_time(self):
        self.send_cmd(CMD_GET_TIME)
        returned = self.get_response()
        if returned[0] == STATUS_OK and returned[1]:
            datetime_lst = returned[1].replace('+CIPSNTPTIME:', '').split(' ')
            datetime_lst = [item for item in datetime_lst if item.strip()]  # 去除空元素
            return True, '', datetime_lst
        
        return False, returned[1], []
    
    
    def mqtt_set(self, host, port, client_id, will_topic=''):
        """
        MQTT 设置
        参数host: mqtt主机名
        参数port: mqtt主机端口
        参数client_id: 服务端id
        参数will_topic: 遗嘱主题
        返回: 
        """
        cmd_usercfg = CMD_SET_MQTTUSERCFG.format(client_id=client_id)
        self.send_cmd(cmd_usercfg)
        returned = self.get_response(only_return_status=True)
        
        cmd_config = CMD_SET_MQTTCONNCFG.format(topic=will_topic)
        self.send_cmd(cmd_config)
        returned = self.get_response()
        
        cmd_conn = CMD_SET_MQTTCONN.format(host=host,port=port)
        self.send_cmd(cmd_conn)
        returned = self.get_response()
        
        
    # MQTT 发布消息    
    def mqtt_publish(self,topic,data):  
        cmd_pub = CMD_SET_MQTTPUB.format(topic=topic,data=data)
        self.send_cmd(cmd_pub)
        returned = self.get_response(only_return_status=True)
        return returned

    # MQTT 订阅消息    
    def mqtt_subscribe(self,topic):  
        cmd_sub = CMD_SET_MQTTSUB.format(topic=topic)
        self.send_cmd(cmd_sub)
        returned = self.get_response(only_return_status=True)
        return returned

    # MQTT 取消订阅    
    def mqtt_unsubscribe(self,topic):  
        cmd_unsub = CMD_SET_MQTTUNSUB.format(topic=topic)
        self.send_cmd(cmd_unsub)
        returned = self.get_response(only_return_status=True)
        return returned
    
    # MQTT 查询broker信息(属性)
    @property
    def mqtt_brokerinfo(self):
        self.send_cmd(CMD_GET_MQTTCONN)
        returned = self.get_response()
        if returned[0] == STATUS_OK and returned[1]:
            broker_str = returned[1].split(':')[1]
            broker_lst = broker_str.split(',')
            
            state = int(broker_lst[1])  # mqtt状态(0-6)
            if state > 0:  # 只有初始化后读取到的值才是非空, 否则下面的int()会报错
                scheme = int(broker_lst[2])  # 
                host = broker_lst[3].replace('"','')  # 服务器地址
                port = int(broker_lst[4].replace('"',''))  # 端口, 字符型
                path = broker_lst[5].replace('"','')  #
                reconnect = int(broker_lst[6])  #重连设置(0,1)
                
                return True, '', state, scheme, host, port, path, reconnect
            else:
                return True, '', state, 0, '', 0, '', 0
            

        return False, returned[1], 0, 0, '', 0, '', 0
    
    # MQTT 查询订阅主题等信息(属性)  注意:订阅了多个主题,每个主题占一行
    @property
    def mqtt_get_subscribed_info(self):
        self.send_cmd(CMD_GET_MQTTSUB)
        returned = self.get_response()
        
        if returned[0] == STATUS_OK and returned[1]:
            subscribed_lst = returned[1].split('\r\n')  # 先分行提取各个主题 
            topics = []  # 存储主题名称的列表
            
            for i, subscribed_str in enumerate(subscribed_lst):
                topic_str = subscribed_str.split(':')[1]  # 去除前面的命令字符串
                topic_lst = topic_str.split(',')
                if i == 0:
                    state = int(topic_lst[1])  # mqtt状态(0-6),只取第一个主题的
                    qos = int(topic_lst[3])  #qos设置(0,1,2),只取第一个主题的
                topics.append(topic_lst[2].replace('"',''))  # 已订阅主题
            
            return True, '', state, topics, qos

        return False, returned[1], 0, [], 0


    # MQTT 读取订阅的消息
    def mqtt_wait_for_subscribed(self, timeout=TIMEOUT_DEFAULT):
        """
        等待接收已订阅的消息.该消息没有结束语,但包含消息长度, 读取到指定长度即表示接收完成
        消息类似: b'+MQTTSUBRECV:0,"test01",1,a\r\n'
        timeout: 超时时间,单位:秒
        返回:
        """
        
        # 初始化
        prvMills = utime.ticks_ms()  # 接收响应的起始时间
        response = b""  # 响应的初始值
        return_status = STATUS_DEFAULT  #返回状态
        received_lst=[]  # 存储接收数据的列表
        is_subscribed_info = 0  # 接收的消息是否是订阅消息, 0初始值, 1是 ,-1否
        msg_header = ('+'+CMD__MQTTSUBRECV+':').encode('utf8')  

        
        # 接收数据
        while True:
            # 超时退出
            if utime.ticks_diff(utime.ticks_ms(), prvMills) >= timeout:  
                return_status = STATUS_TIMEOUT
                break
            
            # 循环接收
            if self.uart.any():
                response += self.uart.readline()
           
                if is_subscribed_info == 0 :  #尚未确定是订阅消息
                    if len(response) < len(msg_header):  # 接收的消息长度不够,继续接收
                        continue               
                    elif response.startswith(msg_header):  # 确定有标识字符
                        is_subscribed_info = 1
                        response = response[len(msg_header):]
                    else:  # 确定没有标识字符, 仅记录日志,不返回
                        is_subscribed_info = -1
                
                
                # 大部分情况是一次就读取完了, 所以两个if要分别执行
                if is_subscribed_info == 1 :  #已经确定是订阅消息,拆解
                    dismantle_lst = response.split(b',')
                    if len(dismantle_lst) >= 4:
                        if dismantle_lst[2].isdigit():  # byte类型也能这样检测, 无需解码
                            msg_length = int(dismantle_lst[2])
                            if len(dismantle_lst[3].replace(b'\r\n',b''))>= msg_length:  # 接收到指定长度,主动退出(注意不能直接strip,因为消息末尾的空格也会去除导致获取长度不对)
                                received_lst = dismantle_lst
                                return_status = STATUS_OK
                                break
        
        # 后处理
        delay = utime.ticks_diff(utime.ticks_ms(), prvMills)  # 计算延迟,单位毫秒
              
        #调试用
#         print("已接收: ", response, return_status, delay)
#         print('--------------------------')
        
        # 解码
        subject, message = '', ''
        try:
            if is_subscribed_info == 1 :
                if not received_lst:
                    received_lst = response.split(b',')
                    
                if len(dismantle_lst) >= 4:
                    subject = received_lst[1].decode('utf-8')
                    message = received_lst[3].decode('utf-8').strip()
                else:
                    return_status = STATUS_NOT_PASS_MUSTER        
            else:
                return_status = STATUS_NOT_PASS_MUSTER
            
        except Exception as e:
            return_status = STATUS_DECODE_ERROR
            
        # 返回
        return return_status, subject, message, delay
        
    
    # MQTT 断开连接
    def mqtt_disconnect(self):
        self.send_cmd(CMD_SET_MQTTCLEAN)
        returned = self.get_response(only_return_status=True)
    
    
    
    # 获取本设备信息
    def sysinfo(self):
        return self.sendCMD_waitResp("AT+GMR\r\n")
    
    # 
    def restore(self):
        return self.sendCMD_waitResp("AT+RESTORE\r\n")
    
    # ping,判断网络是否联通
    def ping(self, doip):
        return self.sendCMD_waitResp('AT+PING="%s"\r\n' % (doip,))
    
    # 断开连接
    def disconnect(self):
        return self.sendCMD_waitResp("AT+CWQAP\r\n")
    
